package tunnel

import (
	"context"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// Tunnel 端口转发隧道接口
type Tunnel interface {
	// Start 启动隧道
	Start(ctx context.Context) error
	// Stop 停止隧道
	Stop() error
	// Status 获取隧道状态
	Status() *TunnelStatus
}

// TunnelType 隧道类型
type TunnelType string

const (
	TunnelTypeTCP  TunnelType = "tcp"
	TunnelTypeUDP  TunnelType = "udp"
	TunnelTypeHTTP TunnelType = "http"
)

// TunnelConfig 隧道配置
type TunnelConfig struct {
	ID          string
	Type        TunnelType
	LocalAddr   string
	RemoteAddr  string
	BindAddr    string
	Description string
}

// TunnelStatus 隧道状态
type TunnelStatus struct {
	ID          string
	Type        TunnelType
	LocalAddr   string
	RemoteAddr  string
	BindAddr    string
	Status      string // running, stopped, error
	Connections int    // 当前连接数
	BytesIn     int64  // 接收字节数
	BytesOut    int64  // 发送字节数
	StartTime   time.Time
	LastError   string
}

// ConnectionInfo 连接信息
type ConnectionInfo struct {
	ID        string
	LocalAddr string
	RemoteAddr string
	StartTime time.Time
	BytesIn   int64
	BytesOut  int64
}

// TCPTunnel TCP隧道实现
type TCPTunnel struct {
	config      *TunnelConfig
	status      *TunnelStatus
	listener    net.Listener
	connections map[string]*ConnectionInfo
	connMutex   sync.RWMutex
	ctx         context.Context
	cancel      context.CancelFunc
	wg          sync.WaitGroup
}

// NewTCPTunnel 创建新的TCP隧道
func NewTCPTunnel(config *TunnelConfig) *TCPTunnel {
	ctx, cancel := context.WithCancel(context.Background())
	
	return &TCPTunnel{
		config:      config,
		connections: make(map[string]*ConnectionInfo),
		ctx:         ctx,
		cancel:      cancel,
		status: &TunnelStatus{
			ID:         config.ID,
			Type:       config.Type,
			LocalAddr:  config.LocalAddr,
			RemoteAddr: config.RemoteAddr,
			BindAddr:   config.BindAddr,
			Status:     "stopped",
			StartTime:  time.Now(),
		},
	}
}

// Start 启动TCP隧道
func (t *TCPTunnel) Start(ctx context.Context) error {
	// 合并上下文
	mergedCtx, cancel := context.WithCancel(context.Background())
	t.cancel = cancel
	
	go func() {
		select {
		case <-ctx.Done():
			cancel()
		case <-mergedCtx.Done():
		}
	}()
	
	t.ctx = mergedCtx
	
	// 监听本地地址
	listener, err := net.Listen("tcp", t.config.BindAddr)
	if err != nil {
		t.status.LastError = err.Error()
		return fmt.Errorf("failed to listen on %s: %v", t.config.BindAddr, err)
	}
	
	t.listener = listener
	t.status.Status = "running"
	
	// 启动接受连接的goroutine
	t.wg.Add(1)
	go t.acceptLoop()
	
	return nil
}

// acceptLoop 接受连接循环
func (t *TCPTunnel) acceptLoop() {
	defer t.wg.Done()
	
	for {
		select {
		case <-t.ctx.Done():
			return
		default:
		}
		
		// 接受连接
		conn, err := t.listener.Accept()
		if err != nil {
			// 检查是否是关闭错误
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				// 临时错误，继续等待
				continue
			}
			
			// 其他错误，停止监听
			t.status.LastError = err.Error()
			return
		}
		
		// 处理连接
		t.wg.Add(1)
		go t.handleConnection(conn)
	}
}

// handleConnection 处理连接
func (t *TCPTunnel) handleConnection(localConn net.Conn) {
	defer t.wg.Done()
	
	// 创建连接信息
	connInfo := &ConnectionInfo{
		ID:        fmt.Sprintf("%s-%d", t.config.ID, time.Now().UnixNano()),
		LocalAddr: localConn.RemoteAddr().String(),
		RemoteAddr: t.config.RemoteAddr,
		StartTime: time.Now(),
	}
	
	// 添加到连接列表
	t.connMutex.Lock()
	t.connections[connInfo.ID] = connInfo
	t.status.Connections = len(t.connections)
	t.connMutex.Unlock()
	
	// 确保连接最终被清理
	defer func() {
		localConn.Close()
		
		t.connMutex.Lock()
		delete(t.connections, connInfo.ID)
		t.status.Connections = len(t.connections)
		t.connMutex.Unlock()
	}()
	
	// 连接到远程地址
	remoteConn, err := net.Dial("tcp", t.config.RemoteAddr)
	if err != nil {
		t.status.LastError = fmt.Sprintf("failed to connect to remote %s: %v", t.config.RemoteAddr, err)
		return
	}
	defer remoteConn.Close()
	
	// 启动双向数据传输
	t.wg.Add(2)
	go t.transferData(localConn, remoteConn, connInfo, true)
	go t.transferData(remoteConn, localConn, connInfo, false)
	
	// 等待上下文结束或连接关闭
	select {
	case <-t.ctx.Done():
	case <-func() chan struct{} {
		closed := make(chan struct{})
		go func() {
			defer close(closed)
			// 等待连接关闭（简单实现）
			buf := make([]byte, 1)
			localConn.Read(buf)
		}()
		return closed
	}():
	}
}

// transferData 传输数据
func (t *TCPTunnel) transferData(src, dst net.Conn, connInfo *ConnectionInfo, isOutbound bool) {
	defer t.wg.Done()
	
	// 缓冲区
	buffer := make([]byte, 32*1024) // 32KB缓冲区
	
	for {
		select {
		case <-t.ctx.Done():
			return
		default:
		}
		
		// 设置读取超时
		src.SetReadDeadline(time.Now().Add(30 * time.Second))
		
		// 从源读取数据
		n, err := src.Read(buffer)
		if err != nil {
			if err != io.EOF {
				// 检查是否是超时错误
				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
					// 超时继续
					continue
				}
			}
			// 连接关闭或其他错误
			return
		}
		
		// 写入目标
		_, err = dst.Write(buffer[:n])
		if err != nil {
			return
		}
		
		// 更新统计信息
		t.connMutex.Lock()
		if isOutbound {
			connInfo.BytesOut += int64(n)
			t.status.BytesOut += int64(n)
		} else {
			connInfo.BytesIn += int64(n)
			t.status.BytesIn += int64(n)
		}
		t.connMutex.Unlock()
	}
}

// Stop 停止隧道
func (t *TCPTunnel) Stop() error {
	if t.cancel != nil {
		t.cancel()
	}
	
	if t.listener != nil {
		t.listener.Close()
	}
	
	// 等待所有goroutine完成
	t.wg.Wait()
	
	t.status.Status = "stopped"
	return nil
}

// Status 获取隧道状态
func (t *TCPTunnel) Status() *TunnelStatus {
	t.connMutex.RLock()
	defer t.connMutex.RUnlock()
	
	// 返回状态副本
	status := *t.status
	return &status
}

// ListConnections 列出所有连接
func (t *TCPTunnel) ListConnections() []*ConnectionInfo {
	t.connMutex.RLock()
	defer t.connMutex.RUnlock()
	
	connections := make([]*ConnectionInfo, 0, len(t.connections))
	for _, conn := range t.connections {
		// 返回副本
		connCopy := *conn
		connections = append(connections, &connCopy)
	}
	
	return connections
}

// UDPTunnel UDP隧道实现
type UDPTunnel struct {
	config      *TunnelConfig
	status      *TunnelStatus
	conn        *net.UDPConn
	connections map[string]*ConnectionInfo
	connMutex   sync.RWMutex
	ctx         context.Context
	cancel      context.CancelFunc
	wg          sync.WaitGroup
}

// NewUDPTunnel 创建新的UDP隧道
func NewUDPTunnel(config *TunnelConfig) *UDPTunnel {
	ctx, cancel := context.WithCancel(context.Background())
	
	return &UDPTunnel{
		config:      config,
		connections: make(map[string]*ConnectionInfo),
		ctx:         ctx,
		cancel:      cancel,
		status: &TunnelStatus{
			ID:         config.ID,
			Type:       config.Type,
			LocalAddr:  config.LocalAddr,
			RemoteAddr: config.RemoteAddr,
			BindAddr:   config.BindAddr,
			Status:     "stopped",
			StartTime:  time.Now(),
		},
	}
}

// Start 启动UDP隧道
func (t *UDPTunnel) Start(ctx context.Context) error {
	// 合并上下文
	mergedCtx, cancel := context.WithCancel(context.Background())
	t.cancel = cancel
	
	go func() {
		select {
		case <-ctx.Done():
			cancel()
		case <-mergedCtx.Done():
		}
	}()
	
	t.ctx = mergedCtx
	
	// 解析本地地址
	addr, err := net.ResolveUDPAddr("udp", t.config.BindAddr)
	if err != nil {
		t.status.LastError = err.Error()
		return fmt.Errorf("failed to resolve UDP address %s: %v", t.config.BindAddr, err)
	}
	
	// 监听UDP地址
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		t.status.LastError = err.Error()
		return fmt.Errorf("failed to listen on UDP %s: %v", t.config.BindAddr, err)
	}
	
	t.conn = conn
	t.status.Status = "running"
	
	// 启动处理数据的goroutine
	t.wg.Add(1)
	go t.handleUDPData()
	
	return nil
}

// handleUDPData 处理UDP数据
func (t *UDPTunnel) handleUDPData() {
	defer t.wg.Done()
	
	buffer := make([]byte, 64*1024) // 64KB缓冲区
	
	for {
		select {
		case <-t.ctx.Done():
			return
		default:
		}
		
		// 读取UDP数据
		n, clientAddr, err := t.conn.ReadFromUDP(buffer)
		if err != nil {
			// 检查是否是关闭错误
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				continue
			}
			return
		}
		
		// 创建连接信息（UDP是无连接的，这里用客户端地址标识）
		connID := fmt.Sprintf("%s-%s", t.config.ID, clientAddr.String())
		
		t.connMutex.Lock()
		connInfo, exists := t.connections[connID]
		if !exists {
			connInfo = &ConnectionInfo{
				ID:        connID,
				LocalAddr: clientAddr.String(),
				RemoteAddr: t.config.RemoteAddr,
				StartTime: time.Now(),
			}
			t.connections[connID] = connInfo
			t.status.Connections = len(t.connections)
		}
		connInfo.BytesIn += int64(n)
		t.status.BytesIn += int64(n)
		t.connMutex.Unlock()
		
		// 连接到远程地址并转发数据
		t.wg.Add(1)
		go t.forwardUDPData(buffer[:n], clientAddr, connInfo)
	}
}

// forwardUDPData 转发UDP数据
func (t *UDPTunnel) forwardUDPData(data []byte, clientAddr *net.UDPAddr, connInfo *ConnectionInfo) {
	defer t.wg.Done()
	
	// 连接到远程地址
	remoteAddr, err := net.ResolveUDPAddr("udp", t.config.RemoteAddr)
	if err != nil {
		t.status.LastError = fmt.Sprintf("failed to resolve remote UDP address %s: %v", t.config.RemoteAddr, err)
		return
	}
	
	remoteConn, err := net.DialUDP("udp", nil, remoteAddr)
	if err != nil {
		t.status.LastError = fmt.Sprintf("failed to dial remote UDP address %s: %v", t.config.RemoteAddr, err)
		return
	}
	defer remoteConn.Close()
	
	// 发送数据到远程地址
	_, err = remoteConn.Write(data)
	if err != nil {
		t.status.LastError = fmt.Sprintf("failed to write to remote UDP address %s: %v", t.config.RemoteAddr, err)
		return
	}
	
	connInfo.BytesOut += int64(len(data))
	t.status.BytesOut += int64(len(data))
	
	// 读取响应数据
	response := make([]byte, 64*1024)
	remoteConn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err := remoteConn.Read(response)
	if err != nil {
		// 不记录超时错误
		if netErr, ok := err.(net.Error); !ok || !netErr.Timeout() {
			t.status.LastError = fmt.Sprintf("failed to read from remote UDP address %s: %v", t.config.RemoteAddr, err)
		}
		return
	}
	
	// 将响应数据发送回客户端
	_, err = t.conn.WriteToUDP(response[:n], clientAddr)
	if err != nil {
		t.status.LastError = fmt.Sprintf("failed to write to client %s: %v", clientAddr.String(), err)
		return
	}
	
	connInfo.BytesIn += int64(n)
	t.status.BytesIn += int64(n)
}

// Stop 停止UDP隧道
func (t *UDPTunnel) Stop() error {
	if t.cancel != nil {
		t.cancel()
	}
	
	if t.conn != nil {
		t.conn.Close()
	}
	
	// 等待所有goroutine完成
	t.wg.Wait()
	
	t.status.Status = "stopped"
	return nil
}

// Status 获取UDP隧道状态
func (t *UDPTunnel) Status() *TunnelStatus {
	t.connMutex.RLock()
	defer t.connMutex.RUnlock()
	
	// 返回状态副本
	status := *t.status
	return &status
}

// ListConnections 列出所有UDP"连接"
func (t *UDPTunnel) ListConnections() []*ConnectionInfo {
	t.connMutex.RLock()
	defer t.connMutex.RUnlock()
	
	connections := make([]*ConnectionInfo, 0, len(t.connections))
	for _, conn := range t.connections {
		// 返回副本
		connCopy := *conn
		connections = append(connections, &connCopy)
	}
	
	return connections
}

// TunnelManager 隧道管理器
type TunnelManager struct {
	tunnels map[string]Tunnel
	mutex   sync.RWMutex
}

// NewTunnelManager 创建新的隧道管理器
func NewTunnelManager() *TunnelManager {
	return &TunnelManager{
		tunnels: make(map[string]Tunnel),
	}
}

// AddTunnel 添加隧道
func (tm *TunnelManager) AddTunnel(tunnel Tunnel) error {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()
	
	status := tunnel.Status()
	tm.tunnels[status.ID] = tunnel
	
	return nil
}

// RemoveTunnel 移除隧道
func (tm *TunnelManager) RemoveTunnel(tunnelID string) error {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()
	
	tunnel, exists := tm.tunnels[tunnelID]
	if !exists {
		return fmt.Errorf("tunnel %s not found", tunnelID)
	}
	
	// 停止隧道
	tunnel.Stop()
	
	delete(tm.tunnels, tunnelID)
	return nil
}

// StartTunnel 启动隧道
func (tm *TunnelManager) StartTunnel(tunnelID string, ctx context.Context) error {
	tm.mutex.RLock()
	tunnel, exists := tm.tunnels[tunnelID]
	tm.mutex.RUnlock()
	
	if !exists {
		return fmt.Errorf("tunnel %s not found", tunnelID)
	}
	
	return tunnel.Start(ctx)
}

// StopTunnel 停止隧道
func (tm *TunnelManager) StopTunnel(tunnelID string) error {
	tm.mutex.RLock()
	tunnel, exists := tm.tunnels[tunnelID]
	tm.mutex.RUnlock()
	
	if !exists {
		return fmt.Errorf("tunnel %s not found", tunnelID)
	}
	
	return tunnel.Stop()
}

// GetTunnelStatus 获取隧道状态
func (tm *TunnelManager) GetTunnelStatus(tunnelID string) (*TunnelStatus, error) {
	tm.mutex.RLock()
	tunnel, exists := tm.tunnels[tunnelID]
	tm.mutex.RUnlock()
	
	if !exists {
		return nil, fmt.Errorf("tunnel %s not found", tunnelID)
	}
	
	return tunnel.Status(), nil
}

// ListTunnels 列出所有隧道
func (tm *TunnelManager) ListTunnels() []*TunnelStatus {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()
	
	statuses := make([]*TunnelStatus, 0, len(tm.tunnels))
	for _, tunnel := range tm.tunnels {
		statuses = append(statuses, tunnel.Status())
	}
	
	return statuses
}

// CreateTunnel 创建隧道
func (tm *TunnelManager) CreateTunnel(config *TunnelConfig) (Tunnel, error) {
	var tunnel Tunnel
	
	switch config.Type {
	case TunnelTypeTCP:
		tunnel = NewTCPTunnel(config)
	case TunnelTypeUDP:
		tunnel = NewUDPTunnel(config)
	default:
		return nil, fmt.Errorf("unsupported tunnel type: %s", config.Type)
	}
	
	tm.mutex.Lock()
	tm.tunnels[config.ID] = tunnel
	tm.mutex.Unlock()
	
	return tunnel, nil
}